"""
示例: https://tensorflow.google.cn/tutorials/keras/text_classification_with_hub?hl=zh-cn
本教程演示了使用 Tensorflow Hub 和 Keras 进行迁移学习的基本应用。
"""
import os

import tensorflow as tf
import tensorflow_datasets as tfds
import tensorflow_hub as hub

os.environ['TFDS_DATA_DIR'] = r'E:\tensorflow\tensorflow_learn_code\datasets'
print("Version: ", tf.__version__)
print("Eager mode: ", tf.executing_eagerly())
print("Hub version: ", hub.__version__)
print("GPU is", "available" if tf.config.list_physical_devices("GPU") else "NOT AVAILABLE")

# 下载 IMDB 数据集
# Split the training set into 60% and 40% to end up with 15,000 examples
# for training, 10,000 examples for validation and 25,000 examples for testing.
train_data, validation_data, test_data = tfds.load(
    name="imdb_reviews",
    split=('train[:60%]', 'train[60%:]', 'test'),
    as_supervised=True)

# 探索数据
"""
。每个样本都是一个代表电影评论的句子和一个相应的标签。句子未经过任何预处理。
标签是一个整数值（0 或 1），其中 0 表示负面评价，1 表示正面评价。
"""
train_examples_batch, train_labels_batch = next(iter(train_data.batch(10)))
# 我们来打印下前十个样本
# print(f"examples:{train_examples_batch}")
# 我们再打印下前十个标签。
# print(f"label:{train_labels_batch}")

embedding = "https://tfhub.dev/google/nnlm-en-dim50/2"
"""
加载一个预训练的 TensorFlow Hub 模型，并将其封装为一个 Keras 层。这样可以在 Keras 模型中使用这个
预训练模型，例如作为特征提取器或作为模型的一部分进行微调。
"""
hub_layer = hub.KerasLayer(embedding, input_shape=[],
                           dtype=tf.string, trainable=True)
hub_layer(train_examples_batch[:3])  # 提取后三个样本
print(hub_layer)

# 构建模型
# 定义输入层
input_layer = tf.keras.layers.Input(shape=[], dtype=tf.string)


# 自定义层来包装hub_layer
class HubLayerWrapper(tf.keras.layers.Layer):
    """
    初始化HubLayerWrapper层
    :param hub_layer: 传入的Hub层，用于包装
    :param kwargs: 其他关键字参数
    """

    def __init__(self, hub_layer, **kwargs):
        super(HubLayerWrapper, self).__init__(**kwargs)
        self.hub_layer = hub_layer

    """
    构建前向传播过程
    :param inputs: 输入数据
    :return: 通过hub_layer处理后的输入数据
    """

    def call(self, inputs):
        return self.hub_layer(inputs)


# 使用自定义层
x = HubLayerWrapper(hub_layer)(input_layer)  # 自动调用call方法，input_layer作为call方法的参数
x = tf.keras.layers.Dense(16, activation='relu')(x)  # 添加一个全连接层，使用ReLU激活函数
output_layer = tf.keras.layers.Dense(1)(x)  # 输出层，使用线性激活函数

# 构建模型
model = tf.keras.Model(inputs=input_layer, outputs=output_layer)
model.summary()  # 打印模型摘要

model.compile(optimizer='adam',
              loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
              metrics=['accuracy'])

# 训练模型并记录训练过程
# 使用训练数据进行训练，每次训练 epochs 次数，同时在验证数据上进行验证
# shuffle(10000) 用于在训练前对数据进行随机打乱，以改善训练效果
# batch(512) 指定每次训练使用的样本数量为 512
# validation_data.batch(512) 指定验证数据集，同样以 512 的批次进行验证
# verbose=1 在训练过程中输出进度条
history = model.fit(train_data.shuffle(10000).batch(512),
                    epochs=10,
                    validation_data=validation_data.batch(512),
                    verbose=1)

results = model.evaluate(test_data.batch(512), verbose=2)

# 遍历模型的指标名称和对应的结果值
for name, value in zip(model.metrics_names, results):
    # 格式化输出每个指标名称和对应的结果值
    print("%s: %.3f" % (name, value))
